package com.bw.test1;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.regression.LassoRegPredictBatchOp;
import com.alibaba.alink.operator.batch.regression.LassoRegTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class LassoRegTrainBatchOpTest {
    @Test
    public void testLassoRegTrainBatchOp() throws Exception {

        BatchOperator.setParallelism(1);
        // 1、数据源
        List<Row> df = Arrays.asList(
                Row.of(2, 1, 1),
                Row.of(3, 2, 1),
                Row.of(4, 3, 2),
                Row.of(2, 4, 1),
                Row.of(2, 2, 1),
                Row.of(4, 3, 2),
                Row.of(1, 2, 1),
                Row.of(5, 3, 3)
        );
        // 2、选择特征和label
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");
        // Lasso-->套索
        BatchOperator<?> lasso = new LassoRegTrainBatchOp()
                .setLambda(0.1)// 惩罚系数
                .setFeatureCols("f0", "f1")
                .setLabelCol("label");
        // 3、训练模型
        BatchOperator model = batchData.link(lasso);

        // 4、预测
        BatchOperator<?> predictor = new LassoRegPredictBatchOp()
                .setPredictionCol("pred");

        predictor.linkFrom(model, batchData).print();
    }
}